import os
import joblib
from sklearn.metrics import classification_report

from utils.data_loader import load_data
from utils.evaluator import evaluate_model
from utils.preprocessor import preprocess_data
from utils.log import setup_logger


def predict_attrition(input_data):
    # 配置路径
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    fig_path = '../fig'
    # 设置日志
    log_path, log_filename = os.path.join('../log', 'predict'), 'predict'
    logger = setup_logger(log_path, log_filename)
    logger.info("开始预测数据...")
    print("开始预测数据...")
    # 预处理输入数据
    input_data, feature_names = preprocess_data(input_data)

    #  划分数据集
    X = input_data.drop(columns=['Attrition'])
    y = input_data['Attrition']

    # 配置模型加载路径
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    model_path = os.path.join(base_dir, 'model', 'xgb_model.pkl')

    # 加载模型
    model = joblib.load(model_path)

    # 预测
    predictions = model.predict(X)
    probabilities = model.predict_proba(X)[:, 1]

    # 添加预测结果到原始数据
    input_data['Attrition_Prediction'] = predictions
    input_data['Attrition_Probability'] = probabilities

    # 绘制ROC曲线
    logger.info("正在评估模型...")
    print("正在评估模型...")
    report = classification_report(y, model.predict(X))
    roc_auc = evaluate_model(model, X, y, fig_path, 'xgb', 'predict')
    print(f"评估报表: {report}")
    logger.info(f"模型ROC-AUC得分: {roc_auc:.4f}")
    print(f"模型ROC-AUC得分: {roc_auc:.4f}")

    return input_data


if __name__ == "__main__":
    base_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
    data_path = os.path.join(base_dir, 'data')

    # 加载新数据
    data = load_data(data_path, 'test2.csv')

    # 进行预测
    results = predict_attrition(data)

    # 保存结果
    results.to_csv(os.path.join(base_dir, 'data', 'predictions.csv'), index=False)
    print("预测数据已保存至 data/predictions.csv")
